#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
数据预览程序
从指定路径获取和查看数据，预览每个数据文件的表头
"""

import os
import pandas as pd
import numpy as np
from pathlib import Path
import json
from typing import List, Dict, Any

class DataPreview:
    def __init__(self, data_path: str = "/mnt/contest-data/"):
        """
        初始化数据预览器
        
        Args:
            data_path: 数据文件夹路径
        """
        self.data_path = Path(data_path)
        self.supported_formats = ['.csv', '.xlsx', '.xls', '.json', '.parquet', '.txt', '.tsv']
        
    def get_data_files(self) -> List[Path]:
        """
        获取数据文件夹中所有支持的数据文件
        
        Returns:
            数据文件路径列表
        """
        data_files = []
        
        if not self.data_path.exists():
            print(f"警告: 路径 {self.data_path} 不存在")
            return data_files
            
        for file_path in self.data_path.rglob('*'):
            if file_path.is_file() and file_path.suffix.lower() in self.supported_formats:
                data_files.append(file_path)
                
        return sorted(data_files)
    
    def read_file_header(self, file_path: Path, max_rows: int = 5) -> Dict[str, Any]:
        """
        读取文件表头信息
        
        Args:
            file_path: 文件路径
            max_rows: 最大预览行数
            
        Returns:
            包含文件信息的字典
        """
        file_info = {
            'file_path': str(file_path),
            'file_name': file_path.name,
            'file_size': file_path.stat().st_size,
            'file_extension': file_path.suffix.lower(),
            'columns': [],
            'data_types': {},
            'shape': None,
            'preview_data': None,
            'error': None
        }
        
        try:
            # 根据文件类型读取数据
            if file_path.suffix.lower() == '.csv':
                df = pd.read_csv(file_path, nrows=max_rows)
            elif file_path.suffix.lower() in ['.xlsx', '.xls']:
                df = pd.read_excel(file_path, nrows=max_rows)
            elif file_path.suffix.lower() == '.json':
                with open(file_path, 'r', encoding='utf-8') as f:
                    data = json.load(f)
                if isinstance(data, list) and len(data) > 0:
                    df = pd.DataFrame(data[:max_rows])
                else:
                    df = pd.DataFrame([data])
            elif file_path.suffix.lower() == '.parquet':
                df = pd.read_parquet(file_path)
                df = df.head(max_rows)
            elif file_path.suffix.lower() in ['.txt', '.tsv']:
                # 尝试不同的分隔符
                separators = ['\t', ',', ';', '|']
                df = None
                for sep in separators:
                    try:
                        df = pd.read_csv(file_path, sep=sep, nrows=max_rows)
                        if len(df.columns) > 1:  # 如果成功分割出多列，则使用此分隔符
                            break
                    except:
                        continue
                if df is None:
                    df = pd.read_csv(file_path, nrows=max_rows)
            else:
                raise ValueError(f"不支持的文件格式: {file_path.suffix}")
            
            # 获取完整数据信息（不限制行数）
            try:
                if file_path.suffix.lower() == '.csv':
                    full_df = pd.read_csv(file_path)
                elif file_path.suffix.lower() in ['.xlsx', '.xls']:
                    full_df = pd.read_excel(file_path)
                elif file_path.suffix.lower() == '.parquet':
                    full_df = pd.read_parquet(file_path)
                else:
                    full_df = df
                    
                file_info['shape'] = full_df.shape
            except:
                file_info['shape'] = df.shape
            
            # 填充文件信息
            file_info['columns'] = df.columns.tolist()
            file_info['data_types'] = df.dtypes.astype(str).to_dict()
            file_info['preview_data'] = df.to_dict('records')
            
        except Exception as e:
            file_info['error'] = str(e)
            
        return file_info
    
    def format_file_size(self, size_bytes: int) -> str:
        """
        格式化文件大小
        
        Args:
            size_bytes: 文件大小（字节）
            
        Returns:
            格式化后的文件大小字符串
        """
        if size_bytes < 1024:
            return f"{size_bytes} B"
        elif size_bytes < 1024**2:
            return f"{size_bytes/1024:.1f} KB"
        elif size_bytes < 1024**3:
            return f"{size_bytes/(1024**2):.1f} MB"
        else:
            return f"{size_bytes/(1024**3):.1f} GB"
    
    def preview_all_data(self, max_rows: int = 5) -> List[Dict[str, Any]]:
        """
        预览所有数据文件
        
        Args:
            max_rows: 每个文件最大预览行数
            
        Returns:
            所有文件信息列表
        """
        data_files = self.get_data_files()
        
        if not data_files:
            print(f"在路径 {self.data_path} 中未找到支持的数据文件")
            return []
        
        print(f"找到 {len(data_files)} 个数据文件:")
        print("=" * 80)
        
        all_file_info = []
        
        for i, file_path in enumerate(data_files, 1):
            print(f"\n[{i}/{len(data_files)}] 处理文件: {file_path.name}")
            
            file_info = self.read_file_header(file_path, max_rows)
            all_file_info.append(file_info)
            
            # 打印文件信息
            if file_info['error']:
                print(f"❌ 读取失败: {file_info['error']}")
            else:
                print(f"📁 文件路径: {file_info['file_path']}")
                print(f"📊 文件大小: {self.format_file_size(file_info['file_size'])}")
                print(f"📋 数据形状: {file_info['shape']} (行数, 列数)")
                print(f"🏷️  列名 ({len(file_info['columns'])}个):")
                
                # 打印列名和数据类型
                for col in file_info['columns']:
                    dtype = file_info['data_types'].get(col, 'unknown')
                    print(f"   - {col} ({dtype})")
                
                # 打印预览数据
                if file_info['preview_data']:
                    print(f"\n👀 数据预览 (前{len(file_info['preview_data'])}行):")
                    preview_df = pd.DataFrame(file_info['preview_data'])
                    print(preview_df.to_string(index=False, max_cols=10, max_colwidth=20))
            
            print("-" * 80)
        
        return all_file_info
    
    def save_summary_report(self, file_info_list: List[Dict[str, Any]], output_path: str = "data_summary.json"):
        """
        保存数据摘要报告
        
        Args:
            file_info_list: 文件信息列表
            output_path: 输出文件路径
        """
        summary = {
            'total_files': len(file_info_list),
            'scan_path': str(self.data_path),
            'successful_files': len([f for f in file_info_list if not f['error']]),
            'failed_files': len([f for f in file_info_list if f['error']]),
            'file_details': file_info_list
        }
        
        with open(output_path, 'w', encoding='utf-8') as f:
            json.dump(summary, f, ensure_ascii=False, indent=2)
        
        print(f"\n📄 数据摘要报告已保存到: {output_path}")

def main():
    """
    主函数
    """
    print("🔍 数据预览程序启动")
    print("=" * 80)
    
    # 创建数据预览器
    previewer = DataPreview("/mnt/contest-data/")
    
    # 预览所有数据
    file_info_list = previewer.preview_all_data(max_rows=5)
    
    # 保存摘要报告
    if file_info_list:
        previewer.save_summary_report(file_info_list)
    
    print("\n✅ 数据预览完成!")

if __name__ == "__main__":
    main()